#include "kernel_operator.h"
#include <type_traits>
#include <cmath>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename T> class KernelReplicationPad2d {
public:
    __aicore__ inline KernelReplicationPad2d() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR paddings, GM_ADDR y, 
                            int32_t size, int32_t x_ndarray[], int32_t x_dimensional) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->x_dimensional = x_dimensional;

        paddingsGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_PADDINGS *>(paddings), 1);
        size = 1;
        for(int i = 0; i < x_dimensional; i++)
        {
            this->x_ndarray[i] = x_ndarray[i];
            this->y_ndarray[i] = x_ndarray[i];
            size *= x_ndarray[i];
        }
        this->y_ndarray[x_dimensional-2] = this->y_ndarray[x_dimensional-2] + paddingsGm.GetValue(2) + paddingsGm.GetValue(3);
        this->y_ndarray[x_dimensional-1] = this->y_ndarray[x_dimensional-1] + paddingsGm.GetValue(0) + paddingsGm.GetValue(1);
        this->size =size;

        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(x), 1);
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), 1);


        this->loopCount = 1;
        for(int i = 0; i < this->x_dimensional-2; i++)
        {
            this->loopCount *= this->y_ndarray[i];
        }
        this->x_interval[0] = this->x_ndarray[x_dimensional-1];
        this->x_interval[1] = this->x_ndarray[x_dimensional-2];

        this->y_interval[0] = this->y_ndarray[x_dimensional-1];
        this->y_interval[1] = this->y_ndarray[x_dimensional-2];

        if constexpr (std::is_same_v<T, float>)
        {
            this->x_interval_32[0] = (this->x_interval[0]+paddingsGm.GetValue(0)+7)/8*8;
            this->y_interval_32[0] = (this->y_interval[0]+7)/8*8;
        }
        else
        {
            this->x_interval_32[0] = (this->x_interval[0]+paddingsGm.GetValue(0)+15)/16*16;
            this->y_interval_32[0] = (this->y_interval[0]+15)/16*16;
        }

        if constexpr (std::is_same_v<T, float>)
        {
            this->y_st[0] = (paddingsGm.GetValue(0) + this->x_interval[0])/8*8;
            this->y_le[0] = this->y_interval[0] - this->y_st[0];
            this->y_le_32[0] = this->y_le[0] /8*8;

            this->y_st[1] = (paddingsGm.GetValue(0))/8*8;
            this->y_le[1] = this->x_interval[0]+paddingsGm.GetValue(0) - this->y_st[1];
            this->y_le_32[1] = this->y_le[1] /8*8;
        }
        else
        {
            this->y_st[0] = (paddingsGm.GetValue(0) + this->x_interval[0])/16*16;
            this->y_le[0] = this->y_interval[0] - this->y_st[0];
            this->y_le_32[0] = this->y_le[0]/16*16;

            this->y_st[1] = (paddingsGm.GetValue(0))/16*16;
            this->y_le[1] = this->x_interval[0]+paddingsGm.GetValue(0) - this->y_st[1];
            this->y_le_32[1] = this->y_le[1] /16*16;
        }
        // for(int i = 0; i < this->x_dimensional; i++)
        // {
        //     loopCount *= this->x_ndarray[i];
        // }

        // loopCount = loopCount / this->x_ndarray[this->dim];

        // cycles = this->x_ndarray[this->dim];
        // for(int i = this->dim+1; i < this->x_dimensional; i++)
        // {
        //     interval *= this->x_ndarray[i];
        // }

        // this->cycles = cycles;
        // this->interval = interval;
        // this->loopCount = loopCount;

        // if(this->interval == 1)
        // {
            pipe.InitBuffer(inQueueX, BUFFER_NUM, this->x_interval_32[0] * sizeof(DTYPE_X));
            pipe.InitBuffer(outQueueY, BUFFER_NUM, this->y_interval_32[0] * sizeof(DTYPE_Y));
            // pipe.InitBuffer(QueueBuff, 1024 * sizeof(DTYPE_X));
        // }
    }

    __aicore__ inline void Process() {
        
        // int x1 = 0, x2 = 0;
        // for(int i = 0; i < this->loopCount; i++)
        // {
        //     for(int y2=0; y2 < this->y_interval[1]; y2++)
        //     {
        //         if(y2 <= paddingsGm.GetValue(2))
        //         {
        //             x2 = 0;
        //         }
        //         else if(y2 >= paddingsGm.GetValue(2) + this->x_interval[1])
        //         {
        //         }
        //         else
        //         {
        //             x2 ++;
        //         }
        //         for(int y1=0; y1 < this->y_interval[0] ; y1++)
        //         {
        //             if(y1 <= paddingsGm.GetValue(0))
        //             {
        //                 x1 = 0;
        //             }
        //             else if(y1 >= paddingsGm.GetValue(0) + this->x_interval[0])
        //             {
        //             }
        //             else
        //             {
        //                 x1 ++;
        //             }
        //             yGm.SetValue(i*this->y_interval[1]*this->y_interval[0]  + y2 * this->y_interval[0]  + y1, xGm.GetValue(i*this->x_interval[1]*this->x_interval[0] + x2 * this->x_interval[0] + x1));
        //         }
        //     }
        // }

        int x1 = 0, x2 = 0;
        
        for(int i = 0; i < this->loopCount; i++)
        {
            CopyIn(i*this->x_interval[1]*this->x_interval[0]);
            Compute(0);
            CopyOut(i*this->y_interval[1]*this->y_interval[0], paddingsGm.GetValue(2));
            x2 = 0;
            for(int y2=paddingsGm.GetValue(2); y2 < paddingsGm.GetValue(2) + this->x_interval[1]; y2++)
            {
                // if(y2 <= paddingsGm.GetValue(2))
                // {
                //     x2 = 0;
                // }
                // else if(y2 >= paddingsGm.GetValue(2) + this->x_interval[1])
                // {
                // }
                // else
                {
                    
                }
                CopyIn(i*this->x_interval[1]*this->x_interval[0] + x2 * this->x_interval[0]);
                Compute(0);
                CopyOut(i*this->y_interval[1]*this->y_interval[0] + y2 * this->y_interval[0],1);
                x2 ++;
            }
            CopyIn(i*this->x_interval[1]*this->x_interval[0] + (this->x_interval[1]-1) * this->x_interval[0]);
            Compute(0);
            CopyOut(i*this->y_interval[1]*this->y_interval[0] + (paddingsGm.GetValue(2) + this->x_interval[1]) * this->y_interval[0], paddingsGm.GetValue(3));
        }

        //float
        //this->x_interval[0] 128    % 8 >64 <=128
        //this->x_interval[1] 120    128 128
        //paddingsGm.GetValue(0))>8  <11 !11 !9 paddingsGm.GetValue(0))!%4
        // paddingsGm.GetValue(1))>4 <=8 
        //(paddingsGm.GetValue(0) + paddingsGm.GetValue(1))%4 !=0




        // if ((paddingsGm.GetValue(0) + paddingsGm.GetValue(1))%4 ==0) // (this->x_interval[0] % 8 == 0)&&   +paddingsGm.GetValue(1)
        // {
        //     TBuf<QuePosition::VECCALC> boox;                
        //     pipe.InitBuffer(boox, 9999999 * sizeof(float));
        //     auto booxxxxx = boox.Get<float>();
        //     Add(booxxxxx, booxxxxx, booxxxxx, 9999999);
        //     Add(booxxxxx[1], booxxxxx[2], booxxxxx[3], 9999999);
        // }
        
        // for(int i = 0; i < this->loopCount; i++)
        // {
        //     for(int y2=0; y2 < this->y_interval[1]; y2++)
        //     {
        //         if(y2 <= paddingsGm.GetValue(2))
        //         {
        //             x2 = 0;
        //         }
        //         else if(y2 >= paddingsGm.GetValue(2) + this->x_interval[1])
        //         {
        //         }
        //         else
        //         {
        //             x2 ++;
        //         }
        //         CopyIn(i*this->x_interval[1]*this->x_interval[0] + x2 * this->x_interval[0]);
        //         Compute(0);
        //         CopyOut(i*this->y_interval[1]*this->y_interval[0] + y2 * this->y_interval[0],1);
        //         // for(int y1=0; y1 < this->y_interval[0] ; y1++)
        //         // {
        //         //     if(y1 <= paddingsGm.GetValue(0))
        //         //     {
        //         //         x1 = 0;
        //         //     }
        //         //     else if(y1 >= paddingsGm.GetValue(0) + this->x_interval[0])
        //         //     {
        //         //     }
        //         //     else
        //         //     {
        //         //         x1 ++;
        //         //     }
        //         //     yGm.SetValue(i*this->y_interval[1]*this->y_interval[0] + y2 * this->y_interval[0]  + y1, xGm.GetValue(i*this->x_interval[1]*this->x_interval[0] + x2 * this->x_interval[0] + x1));
        //         // }
        //     }
        // }

    }

private:
    __aicore__ inline void CopyIn(int32_t x_start) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[x_start-paddingsGm.GetValue(0)], this->x_interval_32[0]);
        // DataCopy(xLocal[this->y_st[1]], xGm[x_start-paddingsGm.GetValue(0)+this->y_st[1]], this->y_le_32[1]);
        inQueueX.EnQue(xLocal);
    }
    // __aicore__ inline void CopyOut(int32_t y_start) {
    //     int length[3];
    //     if constexpr (std::is_same_v<T, float>)
    //     {
    //         length[0] = (paddingsGm.GetValue(0)+7)/8*8;
    //         length[1] = this->x_interval_32[0];
    //         length[2] = (paddingsGm.GetValue(1)+7)/8*8;
    //     }
    //     else
    //     {
    //         length[0] = (paddingsGm.GetValue(0)+15)/16*16;
    //         length[1] = this->x_interval_32[0];
    //         length[2] = (paddingsGm.GetValue(1)+15)/16*16;
    //     }

    //     LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
    //     LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

    //     Duplicate(yLocal, (DTYPE_Y)xLocal.GetValue(0), paddingsGm.GetValue(0));
    //     DataCopy(yGm[y_start], yLocal, length[0]);

    //     // DataCopy(yGm[y_start + paddingsGm.GetValue(0)], xLocal, length[1]);

    //     // Duplicate(yLocal, (DTYPE_Y)xLocal.GetValue(this->x_interval[0]-1), paddingsGm.GetValue(2));
    //     // DataCopy(yGm[y_start + paddingsGm.GetValue(0) + this->x_interval[0]], yLocal, length[2]);

    //     inQueueX.FreeTensor(xLocal);
    //     outQueueY.FreeTensor(yLocal);
    // }
    __aicore__ inline void Compute(int32_t progress) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

        // Duplicate(yLocal[this->y_st[0]], (DTYPE_Y)xLocal.GetValue(paddingsGm.GetValue(0)+this->x_interval[0]-1), this->y_le[0]);
        // Duplicate(yLocal, (DTYPE_Y)0, this->x_interval[0]+paddingsGm.GetValue(0));
        // Add(yLocal, yLocal, xLocal, this->x_interval[0]+paddingsGm.GetValue(0));
        // Duplicate(yLocal, (DTYPE_Y)xLocal.GetValue(paddingsGm.GetValue(0)), paddingsGm.GetValue(0));

        Duplicate(yLocal[this->y_st[0]], (DTYPE_Y)xLocal.GetValue(paddingsGm.GetValue(0)+this->x_interval[0]-1), this->y_le[0]);
        Duplicate(yLocal[this->y_st[1]], (DTYPE_Y)0, this->y_le[1]);
        Add(yLocal[this->y_st[1]], yLocal[this->y_st[1]], xLocal[this->y_st[1]], this->y_le[1]);
        Duplicate(yLocal, (DTYPE_Y)xLocal.GetValue(paddingsGm.GetValue(0)), paddingsGm.GetValue(0));

        // for(int y1=paddingsGm.GetValue(0) + this->x_interval[0]; y1 < this->y_interval[0] ; y1++)
        // {
        //     yLocal.SetValue(y1, xLocal.GetValue(paddingsGm.GetValue(0)+this->x_interval[0]-1));
        // }
        
        // int x1 = 0;
        // Duplicate(yLocal, (DTYPE_Y)xLocal.GetValue(0), paddingsGm.GetValue(0));
        // if constexpr (std::is_same_v<T, float>)
        //     Duplicate(yLocal[(this->y_interval[0] - paddingsGm.GetValue(1))/8*8], (DTYPE_Y)xLocal.GetValue(0), this->y_interval_32[0] - paddingsGm.GetValue(1));
        // else

        // for(int y1=0; y1 < this->y_interval[0] ; y1++)
        // {
        //     if(y1 <= paddingsGm.GetValue(0))
        //     {
        //         x1 = 0;
        //         // yLocal.SetValue(y1, xLocal.GetValue(x1));
        //     }
        //     else if(y1 >= paddingsGm.GetValue(0) + this->x_interval[0])
        //     {
        //         // yLocal.SetValue(y1, xLocal.GetValue(x1));
        //     }
        //     else
        //     {
        //         x1 ++;
        //         yLocal.SetValue(y1, xLocal.GetValue(x1));
        //     }
            
        // }
        inQueueX.FreeTensor(xLocal);
        outQueueY.EnQue<DTYPE_Y>(yLocal);
    }
    __aicore__ inline void CopyOut(int32_t y_start, int32_t progress) {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopyExtParams copyParams{1, (uint16_t)(this->y_interval[0] * sizeof(DTYPE_Y)), 0, 0, 0};
        // DataCopyPad(dstGlobal, dstLocal, copyParams);
        if(progress == 1)
            DataCopyPad(yGm[y_start], yLocal, copyParams);
        else
        {
            for(int i=0;i<progress;i++)
            {
                DataCopyPad(yGm[y_start + i * this->y_interval[0]], yLocal, copyParams);
            }
        }

        // if(progress == 1)
        //     DataCopy(yGm[y_start], yLocal, this->y_interval_32[0]);
        // else
        // {
        //     for(int i=0;i<progress;i++)
        //     {
        //         DataCopy(yGm[y_start + i * this->y_interval[0]], yLocal, this->y_interval_32[0]);
        //     }
        // }
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2;
    
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    GlobalTensor<DTYPE_PADDINGS> paddingsGm;

    int32_t size;
    int32_t x_ndarray[10];
    int32_t y_ndarray[10];
    int32_t x_dimensional;

    int32_t loopCount;
    int32_t x_interval[2];
    int32_t y_interval[2];

    int32_t x_interval_32[2];
    int32_t y_interval_32[2];

    int32_t y_st[3];
    int32_t y_le[3];
    int32_t y_le_32[3];

};
extern "C" __global__ __aicore__ void replication_pad2d(GM_ADDR x, GM_ADDR paddings, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    // TODO: user kernel impl
    KernelReplicationPad2d<DTYPE_X> op;
    op.Init(x, paddings, y, 
            tiling_data.size, tiling_data.x_ndarray, tiling_data.x_dimensional);
    op.Process();
}